from collections import defaultdict
import os
import pickle
from numpy.lib.function_base import select
import pandas as pd
from pandas.core.frame import DataFrame
import datetime
import numpy as np


# 股票信息读写类
class StockDataHelper:
    def __init__(self) -> None:
        self.data_by_date = defaultdict(pd.DataFrame)
        self.goverment_loan = pd.DataFrame()
        self.trade_day_list = []
        self.stock_list = []
        self.factor_list = []

    def init_goverment_loan(self):
        path = os.path.join(os.getcwd(), 'data', '10Y国债收益率.xlsx')
        self.goverment_loan = pd.read_excel(path)

    def init_data(self, file_path):
        """
        :params file_path: 绝对路径
        :return dict{}
        """
        try:
            with open(file_path, "rb") as f:
                data = pickle.load(f)
                self.data_by_date = data
                self.trade_day_list = list(data.keys())
                df = data[self.trade_day_list[0]]
                self.stock_list = list(df.index)

        except IOError:
            print(IOError)
            return {}

    def get_start_end_dayData(self, start_date: str, end_date: str):
        start_date = self.find_trade_day(start_date, self.trade_day_list)
        end_date = self.find_trade_day(end_date, self.trade_day_list)
        data = defaultdict(DataFrame)
        for date in self.trade_day_list:
            if date >= start_date and date <= end_date:
                data[date] = self.data_by_date[date]

        return data

    @staticmethod
    def find_trade_day(date: str, trade_day_list):
        """
        :date: "yyyy-mm-dd"
        :return: 返回date往后最近的一个交易日， 如果无，则返回最后一个交易日
        """
        l, r = 0, len(trade_day_list) - 1
        while l < r:
            mid = (l + r) // 2
            if trade_day_list[mid] >= date:
                r = mid
            else:
                l = mid + 1
        return trade_day_list[l]

    @staticmethod
    def max_drawdown(return_data: pd.Series):
        index_j = np.argmax(np.maximum.accumulate(return_data)-return_data)
        index_i = np.argmax(return_data[:index_j])
        max_drawdown = return_data[index_j]-return_data[index_i]
        return max_drawdown

    @staticmethod
    def get_pre_day(now: str, offet):
        pre_date = (datetime.datetime.strptime(now, '%Y-%m-%d') +
                    datetime.timedelta(-offet)).strftime('%Y-%m-%d')
        return pre_date

    @staticmethod
    def reform_data(data):
        """
        :param data: data.pkl的数据，字典{"trade_date":该trade_date所有股票的数据}
        :return: {'stock_id': 该stock_id的股票数据dataframe,...}
        """
        if data is None:
            return {}
        result = {}
        for date, info in data.items():
            info["trade_date"] = date
            for j in range(len(info.index)):
                stock, data = info.index[j], info.iloc[j: j + 1]
                if stock not in result:
                    result[stock] = data
                else:
                    result[stock] = result[stock].append(data)

        return result

    @staticmethod
    def write_pickle(file_path, name, data):
        """
        ：params file_path: 绝对路径
        """
        try:
            file = os.path.join(file_path, name)
            with open(file, "wb") as f:
                pickle.dump(data, f)
                print("done")
        except IOError:
            print("写入文件出错")
        return None

    @staticmethod
    def get_parital_data_by_k(data, k: int = 100):
        """
        :params data: data.pkl的数据，字典{"trade_date":该trade_date所有股票的数据}
        :return 前k条数据
        """
        res = {}
        for date, info in data.items():
            res[date] = info
            k -= 1
            if k == 0:
                break
        return res
